#!/usr/bin/env python3
# encoding=utf-8

import time
import asyncio
import aiohttp


class AsyncResponse:
    def __init__(self, url: str, status: int, headers: dict, text: str):
        self.status: int = status
        self.url: str = url
        self.headers: dict = headers
        self.text: str = text


class AsyncRequest:
    """
    Async Request Class
    """

    def __init__(self, debug: bool = False):
        self.debug: bool = debug

    async def request(self, method, url: str, **kwargs):
        try:
            Helper.log(self, f"Getting page from {url}...")
            if kwargs['proxy_auth']:
                kwargs['proxy_auth'] = aiohttp.BasicAuth(
                    kwargs['proxy_auth']['user'],
                    kwargs['proxy_auth']['password'])
            else:
                del kwargs['proxy_auth']

            start = time.time()
            async with method(url, **kwargs) as response:
                end = time.time()
                Helper.log(self, f"Get page {url} done, used {end - start}.3f seconds")
                start = time.time()
                text = await response.text()
                end = time.time()
                Helper.log(
                    self,
                    "Get text of %s used %.3f seconds" % (url, end - start))
                res = AsyncResponse(response.url, response.status,
                                    response.headers, text)
        except Exception as e:
            Helper.log(self, "Exception has occured: %s" % str(e))
            res = AsyncResponse(url, f"cannot connect to server {url}", None,
                                str(e))
        return res

    async def get(self,
                  url,
                  headers=None,
                  cookies=None,
                  proxy=None,
                  proxy_auth=None,
                  **kwargs):

        async with aiohttp.ClientSession(cookies=cookies) as session:
            return await self.request(
                session.get,
                url,
                headers=headers,
                proxy=proxy,
                proxy_auth=proxy_auth)

    async def post(self,
                   url,
                   headers=None,
                   data=None,
                   cookies=None,
                   proxy=None,
                   proxy_auth=None,
                   **kwargs):
        async with aiohttp.ClientSession(cookies=cookies) as session:
            return await self.request(
                session.post,
                url,
                headers=headers,
                data=data,
                proxy=proxy,
                proxy_auth=proxy_auth)

    async def put(self,
                  url,
                  headers=None,
                  data=None,
                  cookies=None,
                  proxy=None,
                  proxy_auth=None,
                  **kwargs):
        async with aiohttp.ClientSession(cookies=cookies) as session:
            return await self.request(
                session.put,
                url,
                headers=headers,
                data=data,
                proxy=proxy,
                proxy_auth=proxy_auth)

    async def delete(self,
                     url,
                     headers=None,
                     cookies=None,
                     proxy=None,
                     proxy_auth=None,
                     **kwargs):
        async with aiohttp.ClientSession(cookies=cookies) as session:
            return await self.request(
                session.delete,
                url,
                headers=headers,
                proxy=proxy,
                proxy_auth=proxy_auth)


class AsyncRequestPool:
    def __init__(self,
                 max_request: int = 512,
                 callback=None,
                 debug: bool = False):
        """

        :param max_request:  最大并发
        :param callback:  处理异步请求响应并将返回值添加在result列表里
        接受两个参数，request和response
        :param debug:
        """
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        self.sem = asyncio.Semaphore(max_request)
        self.requests: list = []  # 每个请求的参数的列表
        self.result: list = []
        self.pairs: list = []  # request, response和response的处理结果的字典
        self.callback_func = callback
        self.debug: bool = debug

    def append(self, request: dict):
        """
        request的格式，可以通过在request里面附加参数传入callback
        必填项:method,url;headers和data、cookies等根据需求自行填写
        {
            'method': 'GET',
            'url':'https://www.github.com',
            'headers':{'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/67.0.3396.99 Safari/537.36'}
        } or
        {
            'method': 'POST',
            'url':'http://httpbin.org/post',
            'headers':{'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/67.0.3396.99 Safari/537.36'},
            'data':b'data',
            'args':{
                'a':1,
                'b':2
            }
        }
        'args' means external argument of callback, this class will add it to callback if it
        exists in request dict
        """
        if isinstance(request, dict):
            self.requests.append(request)
        elif isinstance(request, list):
            for r in request:
                self.requests.append(r)
        else:
            print("Invalid request type")

    def callback(self, request: dict, response):
        """
        调用传入的callback，没传入callback则不处理
        """
        if self.callback_func:
            return self.callback_func(request, response)
        return response

    async def send(self, request: dict):
        browser = AsyncRequest(debug=self.debug)
        async with self.sem:

            func = {
                "get": browser.get(**request),
                "post": browser.post(**request),
                "put": browser.put(**request),
                "delete": browser.delete(**request)
            }

            response = await func.get(request['method'].lower())

            data = self.callback(request, response)
            self.result.append(data)
            self.pairs.append({
                'request': request,
                'response': response,
                'data': data
            })

    def run(self):
        """
        在创建好pool，添加完请求后需要调用这个函数用于发出请求，调用完成后result就是处理好的请求结果
        """
        start = time.time()
        self.loop.run_until_complete(
            asyncio.wait(map(self.send, self.requests)))
        end = time.time()
        Helper.log(self, f"AioRequestPool total run {(end - start)}.5f seconds")


class Helper:
    @staticmethod
    def log(cls: AsyncRequest, s):
        if cls.debug:
            print(s)
